import tkinter as tk
from tkinter import filedialog
import cv2
import numpy as np
import pyautogui
from PIL import Image, ImageTk
import json
import os

class ImageFinderApp:
    def __init__(self, root):
        self.root = root
        self.root.title("屏幕图像查找器")
        self.config_file = "config.json"
        
        # 四个位置的图片配置
        self.target_images = [
            {"path": "", "thumbnail": None, "button": None, "label": None},
            {"path": "", "thumbnail": None, "button": None, "label": None},
            {"path": "", "thumbnail": None, "button": None, "label": None},
            {"path": "", "thumbnail": None, "button": None, "label": None}
        ]
        
        # 配置选项
        self.show_preview = tk.BooleanVar(value=True)
        self.auto_click = tk.BooleanVar(value=False)
        
        self.create_widgets()
        self.load_config()
    
    def load_config(self):
        if not os.path.exists(self.config_file):
            return
            
        try:
            with open(self.config_file, 'r') as f:
                config = json.load(f)
                for i, path in enumerate(config.get('target_images', [])):
                    if path and os.path.exists(path):
                        self.target_images[i]['path'] = path
                        self.update_thumbnail(i)
                        
                        # 确保UI组件已初始化后再更新缩略图
                        if hasattr(self, 'target_frame'):
                            self.update_thumbnail(i)
        except (FileNotFoundError, json.JSONDecodeError):
            pass
            
    def save_config(self):
        config = {'target_images': [img['path'] for img in self.target_images]}
        with open(self.config_file, 'w') as f:
            json.dump(config, f)
            
    def update_thumbnail(self, index):
        if not self.target_images[index]["path"]:
            return
            
        # 加载并显示缩略图
        img = cv2.imdecode(np.fromfile(self.target_images[index]["path"], dtype=np.uint8), cv2.IMREAD_COLOR)
        img = cv2.resize(img, (120, 120))
        img = cv2.copyMakeBorder(img, 2,2,2,2, cv2.BORDER_CONSTANT, value=(200,200,200))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        # 转换为Tkinter可用的格式
        img_pil = Image.fromarray(img)
        img_tk = ImageTk.PhotoImage(image=img_pil)
        
        # 保存引用并更新标签
        self.target_images[index]["thumbnail"] = img_tk
        if self.target_images[index]["label"]:
            self.target_images[index]["label"].config(image=img_tk)
        
    def create_widgets(self):

        
        # 阈值调节滑块
        self.threshold_label = tk.Label(self.root, text="匹配阈值: 80")
        self.threshold_label.pack()
        
        self.threshold_slider = tk.Scale(self.root, from_=0, to=100, orient=tk.HORIZONTAL,
                                       command=self.update_threshold)
        self.threshold_slider.set(80)
        self.threshold_slider.pack(fill=tk.X, padx=20, pady=5)
        
        # 四个位置的图片配置区域
        self.target_frame = tk.Frame(self.root)
        self.target_frame.pack(pady=10)
        
        for i in range(4):
            frame = tk.Frame(self.target_frame)
            frame.pack(side=tk.LEFT, padx=10)
            
            # 缩略图标签
            self.target_images[i]["label"] = tk.Label(frame, text=f"位置 {i+1}")
            self.target_images[i]["label"].pack()
            
            # 选择图片按钮
            self.target_images[i]["button"] = tk.Button(
                frame, 
                text="选择图片", 
                command=lambda idx=i: self.select_target_image(idx)
            )
            self.target_images[i]["button"].pack()
        
        # 选项框
        tk.Checkbutton(self.root, text="显示预览", variable=self.show_preview).pack()
        tk.Checkbutton(self.root, text="自动点击", variable=self.auto_click).pack()
        
        # 查找按钮
        tk.Button(self.root, text="依次查找并点击", command=self.find_and_click_all).pack(pady=10)
        
        # 状态显示
        self.status_text = tk.Text(self.root, height=5, width=50)
        self.status_text.pack(pady=5)
        
    def select_target_image(self, index):
        path = filedialog.askopenfilename(
            title=f"选择位置 {index+1} 的图片",
            filetypes=[("图片文件", "*.png;*.jpg;*.jpeg")]
        )
        if path:
            self.save_config()
            self.target_images[index]["path"] = path
            
            # 加载并显示缩略图
            img = cv2.imdecode(np.fromfile(path, dtype=np.uint8), cv2.IMREAD_COLOR)
            img = cv2.resize(img, (100, 100))
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            
            # 转换为Tkinter可用的格式
            img_pil = Image.fromarray(img)
            img_tk = ImageTk.PhotoImage(image=img_pil)
            
            # 保存引用并更新标签
            self.target_images[index]["thumbnail"] = img_tk
            self.target_images[index]["label"].config(image=img_tk)
            
            self.status_text.insert(tk.END, f"已设置位置 {index+1} 的图片: {path}\n")
    
    def find_and_click_all(self):
        for i, target in enumerate(self.target_images):
            if not target["path"]:
                self.status_text.insert(tk.END, f"位置 {i+1} 未设置图片，跳过\n")
                continue
                
            self.status_text.insert(tk.END, f"正在查找位置 {i+1} 的图片...\n")
            self.root.update()
            
            try:
                # 读取模板图像
                template = cv2.imdecode(np.fromfile(target["path"], dtype=np.uint8), cv2.IMREAD_COLOR)
                if template is None:
                    raise ValueError("无法读取模板图片")
                    
                # 获取屏幕截图
                screenshot = pyautogui.screenshot()
                screenshot = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)
                
                # 模板匹配
                result = cv2.matchTemplate(screenshot, template, cv2.TM_CCOEFF_NORMED)
                
                # 获取当前阈值(0-1范围)
                threshold = self.threshold_slider.get() / 100
                
                # 查找匹配位置
                min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
                
                # 当阈值为0时，直接使用最大匹配位置
                if threshold == 0 or max_val >= threshold:
                    x, y = max_loc
                    h, w = template.shape[:2]
                    
                    # 显示结果(可选)
                    if self.show_preview.get():
                        cv2.rectangle(screenshot, (x, y), (x + w, y + h), (0, 255, 0), 2)
                        cv2.imshow(f"匹配结果 {i+1}", screenshot)
                        cv2.waitKey(0)
                        cv2.destroyAllWindows()
                    
                    message = f"位置 {i+1}: 找到匹配图像！位置: ({x}, {y}) 相似度: {max_val:.2f}"
                    self.status_text.insert(tk.END, message + "\n")
                    
                    # 自动点击(可选)
                    if self.auto_click.get():
                        pyautogui.click(x + w//2, y + h//2)
                        self.status_text.insert(tk.END, f"已在中心位置({x + w//2}, {y + h//2})点击\n")
                else:
                    message = f"位置 {i+1}: 未找到匹配图像(阈值: {threshold:.2f})"
                    self.status_text.insert(tk.END, message + "\n")
                    
            except Exception as e:
                self.status_text.insert(tk.END, f"位置 {i+1} 查找失败: {str(e)}\n")
    
    def update_threshold(self, value):
        self.threshold_label.config(text=f"匹配阈值: {value}")
    
    def find_image(self):
        if not self.template_path:
            messagebox.showerror("错误", "请先选择模板图片！")
            return
            
        try:
            # 读取模板图像(处理中文路径问题)
            template = cv2.imdecode(np.fromfile(self.template_path, dtype=np.uint8), cv2.IMREAD_COLOR)
            if template is None:
                raise ValueError("无法读取模板图片")
                
            # 获取屏幕截图
            screenshot = pyautogui.screenshot()
            screenshot = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)
            
            # 模板匹配
            result = cv2.matchTemplate(screenshot, template, cv2.TM_CCOEFF_NORMED)
            
            # 获取当前阈值(0-1范围)
            threshold = self.threshold_slider.get() / 100
            
            # 查找匹配位置
            min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
            
            # 当阈值为0时，直接使用最大匹配位置
            if threshold == 0 or max_val >= threshold:
                x, y = max_loc
                
                # 在匹配位置画矩形
                h, w = template.shape[:2]
                cv2.rectangle(screenshot, (x, y), (x + w, y + h), (0, 255, 0), 2)
                
                # 显示结果(可选)
                if self.show_preview.get():
                    cv2.imshow("匹配结果", screenshot)
                    cv2.waitKey(0)
                    cv2.destroyAllWindows()
                
                message = f"找到匹配图像！位置: ({x}, {y})\n相似度: {max_val:.2f}"
                messagebox.showinfo("成功", message)
                self.status_text.insert(tk.END, message + "\n")
                
                # 自动点击(可选)
                if self.auto_click.get():
                    pyautogui.click(x + w//2, y + h//2)
                    self.status_text.insert(tk.END, f"已在中心位置({x + w//2}, {y + h//2})点击\n")
            else:
                message = f"未找到匹配图像(阈值: {threshold:.2f})"
                messagebox.showinfo("未找到", message)
                self.status_text.insert(tk.END, message + "\n")
                
        except Exception as e:
            messagebox.showerror("错误", f"查找失败: {str(e)}")
            self.status_text.insert(tk.END, f"错误: {str(e)}\n")

if __name__ == "__main__":
    root = tk.Tk()
    app = ImageFinderApp(root)
    root.mainloop()